/////////////////////////////////////////////////////////////////////////
//Copyright (C) 2003 Dekang Lin, lindek@cs.ualberta.ca
//
//Permission to use, copy, modify, and distribute this software for any
//purpose is hereby granted without fee, provided that the above
//copyright notice appear in all copies and that both that copyright
//notice and this permission notice appear in supporting documentation.
//No representations about the suitability of this software for any
//purpose is made. It is provided "as is" without express or implied
//warranty.
//
/////////////////////////////////////////////////////////////////////////

#include <fstream>
#include <iostream>
#include <cmath>
#include <cstdlib>
#include <algorithm>

#include "hmm.h"
#include "logprobs.h"
#include "tables.h"

Transition::Transition(HmmNode* from, HmmNode* to, unsigned long obs)
{
    _from = from;
    _to = to;
    _obs = obs;
    if (_from && _to) {
        _from->outs().push_back(this);
        _to->ins().push_back(this);
    }
}

Hmm::Hmm()
{
    _minLogProb = log(0.000001);
}

double Hmm::getTransProb(Transition* trans)
{
    return _transition.get(trans->_to->state(), trans->_from->state());
}

double Hmm::getEmitProb(Transition* trans)
{
    return _emission.get(trans->_obs, trans->_to->state());
}

void Hmm::forward()
{
    // compute forward probabilities at time 0
    TimeSlot* t0 = _timeSlots[0];
    HmmNode* init = (*t0)[0];
    init->logAlpha(0);

    // compute forward probabilities at time t using the alpha values for time t-1
    for (unsigned int t = 1; t<_timeSlots.size(); t++) {
        TimeSlot* ts = _timeSlots[t];
        for (TimeSlot::iterator it = ts->begin(); it!=ts->end(); it++) {
            vector<Transition*>& ins = (*it)->ins();
            vector<double> logProbs(ins.size());
            for (unsigned int i = 0; i<ins.size(); i++) {
                Transition* trans = ins[i];
                double logProb = trans->_from->logAlpha()+getTransProb(trans)+getEmitProb(trans);
                logProbs[i] = logProb;
            }
            (*it)->logAlpha(sumLogProb(logProbs));
        }
    }
}

double Hmm::viterbi(vector<Transition*>& path)
{
    // set nodes at time 0 according to initial probabilities.
    TimeSlot* ts = _timeSlots[0];
    HmmNode* init = (*ts)[0];
    init->logAlpha(0);

    // find the best path up to path t;
    for (unsigned int t = 1; t<_timeSlots.size(); t++) {//for each seq
        ts = _timeSlots[t];
        for (TimeSlot::iterator it = ts->begin(); it!=ts->end(); it++) {//for each state
            HmmNode* node = *it;
            vector<Transition*>& ins = node->ins();
            double maxProb = log(0.0);
            Transition* bestTrans = 0;
            for (unsigned int i = 0; i<ins.size(); i++) {//for pre state
                Transition* trans = ins[i];
                double logProb = trans->_from->logAlpha()+getTransProb(trans)+getEmitProb(trans);
                if (bestTrans==0 || maxProb<logProb) {
                    bestTrans = trans;
                    maxProb = logProb;
                }
            }

            node->logAlpha(maxProb); // store the highest probability in logAlpha
            node->psi(bestTrans); // store the best transition in psi
        }
    }
    // Find the best node at time T. It will be the last node in the best path
    ts = _timeSlots[_timeSlots.size()-1];
    HmmNode* best = 0;
    for (TimeSlot::iterator it = ts->begin(); it!=ts->end(); it++) {
        HmmNode* node = *it;
        if (best==0 || best->logAlpha()<node->logAlpha())
          best = node;
    }

    // retrieve the nodes in the best path
    for (HmmNode* nd = best; nd;) {
        if (nd->psi()) {
            path.push_back(nd->psi());
            nd = nd->psi()->_from;
        }
        else
          nd = 0;
    } 

    // reverse the path
    for (int i = 0, j=path.size()-1; i<j; i++, j--) {
        Transition* tmp = path[i];
        path[i] = path[j];
        path[j] = tmp;
    }
    // Done
    return best->logAlpha();
}


double Hmm::viterbi2(vector<unsigned>& path, vector<unsigned int> seq)
{
    size_t stateCnt = _emission.size();
    size_t stepCnt = seq.size();
 
    vector <unsigned long> stateList;
    TwoDTable::iterator i = _emission.begin();
    //cout<<stateCnt<<endl;
    while(i != _emission.end())
    {
        stateList.push_back(i->first);
        i++;
    }

    vector< vector<double> > dp(stepCnt, vector<double> (stateCnt, -10000000));
    vector< vector<size_t> > dpS(stepCnt, vector<size_t> (stateCnt, 0));

    //init
    for( size_t i = 0 ; i < dp[0].size() ; i++ )
    {
        dp[0][i] = _transition.get( stateList[i], _initState) + _emission.get(seq[0], stateList[i]);
        //cout<<getStr(stateList[i])<<endl;
        //cout<<_emission.get(seq[0], stateList[i])<<endl;;
        //cout<<_transition.get(stateList[i], _initState)<<endl;;
        //cout<<exp(dp[0][i])<<endl;
    }
    //cout<<"good"<<endl;
    for( size_t i = 1 ; i < seq.size() ; i++ )//step
    {
       for( size_t j = 0 ; j < stateCnt ; j++ )//cur state
       {
           for( size_t k = 0 ; k < stateCnt ; k++ )//per state
           {
               double tempP = dp[i-1][k] + _transition.get(stateList[j], stateList[k]) + _emission.get(seq[i], stateList[j]);
               if( dp[i][j] < tempP )
               {
                   dp[i][j] = tempP;
                   dpS[i][j] = k;
               }
           }
       }
    }
    size_t maxState = 0;
    double maxP = dp[stepCnt-1][0];
    for( size_t i = 0 ; i < stateCnt ; i++ )
    {
        if( maxP < dp[stepCnt-1][i] )
        {
            maxP = dp[stepCnt-1][i];
            maxState = i;
        }
    }
    ///cout<<"maxState "<<maxState<<endl;
    for( int i = stepCnt-1 ; i >= 0 ; i-- )
    {
        path.push_back(stateList[maxState]);
        maxState = dpS[i][maxState];
    }
    reverse(path.begin(), path.end());
    return 0;
}

void Hmm::backward()
{
    int T = _timeSlots.size()-1;
    if (T<1) // no observation
      return;
    for (int t = T; t>=0; t--) {
        TimeSlot* ts = _timeSlots[t];
        for (TimeSlot::iterator it = ts->begin(); it!=ts->end(); it++) {
            HmmNode* node = *it;
            if (t==T)
              node->logBeta(0);
            else {
                vector<Transition*>& outs = node->outs();
                vector<double> logProbs(outs.size());
                for (unsigned int i = 0; i<outs.size(); i++) {
                    Transition* trans = outs[i];
                    double logProb = trans->_to->logBeta()+getTransProb(trans)+getEmitProb(trans);
                    logProbs[i] =logProb;
                }
                node->logBeta(sumLogProb(logProbs));
            }
        }
    }
}

// Compute P(e_1:T) = sum_s P(e_1:T, x_T=s) = sum_s alpha_s(T);
double Hmm::obsProb()
{
    if (_timeSlots.size()<1)
      return 1; // no observations

    forward();
    TimeSlot* last = _timeSlots[_timeSlots.size()-1];
    vector<double> alphaT;
    for (TimeSlot::iterator it = last->begin(); it!=last->end(); it++) {
        alphaT.push_back((*it)->logAlpha());
    }
    return sumLogProb(alphaT);  
}

double Hmm::getPseudoCounts(PseudoCounts& counts)
{
    double PofObs = obsProb(); // this call includes a forward() call.
    backward();
    //  print();

    // Compute the pseudo counts of transitions, emissions, and initializations
    for (unsigned int t = 0; t<_timeSlots.size(); t++) {
        TimeSlot* ts = _timeSlots[t];
        TimeSlot::iterator it = ts->begin();

        // P(X_t=s|e_1:T) = alpha_s(t)*beta_s(t)/P(e_t+1:T|e_1:t)
        // The value sum below is log P(e_t+1:T|e_1:t)
        vector<double> logprobs;
        for (; it!=ts->end(); it++) {
            logprobs.push_back((*it)->logAlpha()+(*it)->logBeta());
        }
        double sum = sumLogProb(logprobs);

        // add the pseudo counts into counts
        for (it = ts->begin(); it!=ts->end(); it++) {
            HmmNode* node = *it;

            //stateCount=P(X_t=s|e_1:T) 
            double stateCount = node->logAlpha()+node->logBeta()-sum; 

            counts.stateCount().add(node->state(), stateCount);
            vector<Transition*>& ins = node->ins();
            unsigned int k;
            for (k = 0; k<ins.size(); k++) {
                Transition* trans = ins[k];
                HmmNode* from = trans->_from;
                double transCount = from->logAlpha()+getTransProb(trans)
                    +getEmitProb(trans)+node->logBeta()-PofObs;
                //	cerr << _str2id.getStr(node->state()) << '\t' 
                //	     << _str2id.getStr(trans->_obs) << '\t'
                //	     << exp(transCount) << endl;
                counts.emitCount().add(node->state(), trans->_obs, transCount);
            }
            vector<Transition*>& outs = node->outs();
            for (k = 0; k<outs.size(); k++) {
                Transition* trans = outs[k];
                HmmNode* to = trans->_to;
                double transCount = node->logAlpha()+getTransProb(trans)
                    +getEmitProb(trans)+to->logBeta()-PofObs;
                counts.transCount().add(node->state(), to->state(), transCount);
            }
        }
    }
    //  counts.print(_str2id);
    return PofObs;
}

void Hmm::baumWelch(vector<vector<unsigned long>*>& sequences, int iterations)
{
    cerr << "Training with Baum-Welch for up to " << iterations << " iterations, using "
        << sequences.size() << " sequences." << endl;
    double prevTotalLogProb = 0;
    for (int k = 0; k<iterations; k++) {
        PseudoCounts counts;
        double totalLogProb = 0;
        for (unsigned int i=0; i<sequences.size(); i++){
            vector<unsigned long>& seq = *sequences[i];
            for (unsigned int j=0; j<seq.size(); j++) {
                addObservation(seq[j]);
            }
            // accumulate the pseudo counts
            totalLogProb += getPseudoCounts(counts);
            reset();
            if ((i+1)%1000==0)
              cerr << "Processed " << i+1 << " sequences" << endl;
        }
        cerr << "Iteration " << k+1 << ' ' << "totalLogProb=" << totalLogProb << endl;
        if (prevTotalLogProb!=0 && (totalLogProb - prevTotalLogProb<1))
          break;
        else
          prevTotalLogProb = totalLogProb;
        updateProbs(counts);
    }
}

void Hmm::trainWithTag(vector<vector<pair<unsigned long, unsigned long> > *> sequences)
{
    PseudoCounts counts;
    _initState = getId("INIT");
    unsigned finalState = getId("FINAL");
    for( unsigned int i = 0 ; i < sequences.size() ; i++ )
    {
        for( unsigned int j = 0 ; j <  sequences[i]->size(); j++ )
        {
            if( j == 0 )
            {
                counts.transCount().add(_initState, (*sequences[i])[j].second, 1);
                counts.stateCount().add(_initState, 1);
            }
            else
            {
                counts.transCount().add((*sequences[i])[j-1].second, (*sequences[i])[j].second, 1);
            }
            if( j == sequences[i]->size() - 1)
            {
                counts.transCount().add((*sequences[i])[j].second, finalState, 1);
                counts.stateCount().add(finalState, 1);
            }
            counts.emitCount().add((*sequences[i])[j].second, (*sequences[i])[j].first, 1);
            counts.stateCount().add((*sequences[i])[j].second, 1);
        }
    }
    updateProbs(counts);
}

void Hmm::updateProbs(PseudoCounts& counts)
{
    _transition.clear();
    _emission.clear();
    for (TwoDTable::iterator i = counts.transCount().begin(); i!=counts.transCount().end(); i++) {
        unsigned long from = i->first;
        double fromCount = counts.stateCount().get(from);
        OneDTable& cnts = *i->second;
        for (OneDTable::iterator j = cnts.begin(); j!=cnts.end(); j++) {
            //      if (j->second-fromCount>=_minLogProb)
            _transition.add(from, j->first, j->second-fromCount);
        }
    }
    for (TwoDTable::iterator s = counts.emitCount().begin(); s!=counts.emitCount().end(); s++) {
        unsigned long state = s->first;
        double stateCount = counts.stateCount().get(state);
        OneDTable& cnts = *s->second;
        for (OneDTable::iterator o = cnts.begin(); o!=cnts.end(); o++) {
            //      if (o->second-stateCount>_minLogProb)
            _emission.add(state, o->first, o->second-stateCount);
        }
    }
    //  saveProbs("");
}

void HmmNode::print()
{
    cerr << _hmm->getStr(state()) << '\t'
        << "alpha=" << exp(logAlpha()) << '\t'
        << "beta=" << exp(logBeta());
    unsigned int i;
    cerr << " (in";
    for (i = 0; i<_ins.size(); i++) {
        cerr << ' ' << _hmm->getStr(_ins[i]->_from->state());
    }
    cerr << ") ";
    cerr << " (out";
    for (i = 0; i<_outs.size(); i++) {
        cerr << ' ' << _hmm->getStr(_outs[i]->_to->state());
    }
    cerr << ")";
    cerr << endl;
}

HmmNode::~HmmNode()
{
    for (unsigned int i = 0; i<_outs.size(); i++) {
        delete _outs[i];
    }
}

TimeSlot::~TimeSlot()
{
    for (TimeSlot::iterator it = begin(); it!=end(); it++) {
        delete (*it);
    }
}

void Hmm::reset()
{
    for (unsigned int t = 0; t<_timeSlots.size(); t++){
        delete _timeSlots[t];
    }
    _timeSlots.clear();
}

Hmm::~Hmm()
{
    reset();
}

void Hmm::addObservation(unsigned long o)
{
    vector<unsigned long> stateIds;
    ULSet* cntx = _emission.getCntx(o);
    if (cntx==0) {
        for (TwoDTable::iterator it = _emission.begin(); it!=_emission.end(); it++) {
            stateIds.push_back(it->first);
        }
    }
    else {
        for (ULSet::iterator it = cntx->begin(); it!=cntx->end(); it++) {
            stateIds.push_back(*it);
        }
    }

    if (_timeSlots.empty()) {
        // create a special state for time 0;
        TimeSlot* t0 = new TimeSlot;
        t0->push_back(new HmmNode(0, _initState, this));
        _timeSlots.push_back(t0);
    }

    TimeSlot* ts = new TimeSlot;
    int time = _timeSlots.size();
    for (unsigned int i = 0; i<stateIds.size(); i++) {
        HmmNode* node = new HmmNode(time, stateIds[i], this);
        ts->push_back(node);
        TimeSlot* prev = _timeSlots[time-1];
        for (TimeSlot::iterator it = prev->begin(); it!=prev->end();it++) {
            ULSet* possibleSrc = _transition.getCntx(node->state());
            if (possibleSrc && possibleSrc->find((*it)->state())!=possibleSrc->end())
              new Transition(*it, node, o);
        }
    }
    _timeSlots.push_back(ts);
}

void Hmm::loadProbs(string name)
{
    string s = name+".trans";
    ifstream transProb(s.c_str());
    string initState;
    transProb >> initState;
    _initState = _str2id.getId(initState);
    _transition.load(transProb, _str2id);

    s = name+".emit";
    ifstream emitProb(s.c_str());
    _emission.load(emitProb, _str2id);
}

void Hmm::saveProbs(string name)
{
    if (name=="") {
        cerr << "transition probabilities:" << endl;
        _transition.save(cerr, _str2id);
        cerr << "---------------------------------" << endl;
        cerr << "emission probabilities:" << endl;
        _emission.save(cerr, _str2id);
        cerr << "---------------------------------" << endl;
    }
    else {
        string s = name+".trans";
        ofstream transProb(s.c_str());
        transProb << _str2id.getStr(_initState) << endl;
        _transition.save(transProb, _str2id);

        s = name+".emit";
        ofstream emitProb(s.c_str());
        _emission.save(emitProb, _str2id);
    }
}

void Hmm::print()
{
    for (unsigned int i = 0; i<_timeSlots.size(); i++) {
        TimeSlot* ts = _timeSlots[i];
        cerr << "TIME=" << i << endl;
        for (unsigned int s = 0; s<ts->size(); s++) {
            (*ts)[s]->print();
        }
    }
}

void Hmm::readSeqs(istream& istrm, vector<vector<unsigned long>*>& sequences)
{
    string line;
    const string delims(" ");
    while (getline(istrm, line)) {
        vector<unsigned long>* seq = new vector<unsigned long>;
        string::size_type begIdx, endIdx;
        begIdx = line.find_first_not_of(delims);
        while (begIdx!=string::npos) {
            if (line[begIdx]=='#') // the rest of the line are comments
              break;
            endIdx = line.find_first_of(delims, begIdx);
            if (endIdx==string::npos) {
                endIdx = line.length();
            }
            string word = line.substr(begIdx, endIdx-begIdx);
            seq->push_back(getId(word));
            begIdx = line.find_first_not_of(delims, endIdx);
        }
        if (seq->size()>0)
          sequences.push_back(seq);
        else
          delete seq;
    }
}


void Hmm::readSeqsWithTag(istream & istrm, vector<vector<pair<unsigned long, unsigned long> > * > &sequences, bool delDate)
{
    string line;
    const string delims(" ");
    const string subDelims("/");
    while (getline(istrm, line)) 
    {
        vector<pair<unsigned long, unsigned long> >* seq = new vector<pair<unsigned long, unsigned long> >;
        string::size_type begIdx, endIdx;
        if( delDate )
        {
            string::size_type firstspace = line.find(" ");//delete the date in the head
            line = line.substr(firstspace + 1, line.size() - firstspace -1 );
        }
        begIdx = line.find_first_not_of(delims);
        while (begIdx!=string::npos) 
        {
            if (line[begIdx]=='#') // the rest of the line are comments
              break;
            endIdx = line.find_first_of(delims, begIdx);
            if (endIdx==string::npos) 
            {
                endIdx = line.length();
            }
            string word = line.substr(begIdx, endIdx-begIdx);
            string::size_type midIdx = word.find(subDelims);
            if( midIdx == string::npos )
              break;
            string wordtmp = word.substr(0, midIdx);
            string tagtmp = word.substr(midIdx+1, word.size() - midIdx - 1);
            if(wordtmp.size())
            {
                seq->push_back(make_pair(getId(wordtmp),getId(tagtmp)));
            }
            //cout<<word.substr(0, midIdx)<<endl;
            //cout<<word.substr(midIdx+1, word.size() - midIdx - 1)<<endl;
            begIdx = line.find_first_not_of(delims, endIdx);
        }
        if (seq->size() > 0)
          sequences.push_back(seq);
        else
          delete seq;
    }
}


void Hmm::preProcess(istream & istrm, vector<vector<pair<unsigned long, unsigned long> > * > &sequences, bool delDate)
{
    readSeqsWithTag(istrm,sequences, delDate);
    for( unsigned int i = 0 ; i < sequences.size() ; i++ )
    {
        unsigned int j = 0;
        while(j < sequences[i]->size())
        {
            string word = getStr((*sequences[i])[j].first);
            string tag = getStr((*sequences[i])[j].second);
            if( word[0] == '[' )
            {
                unsigned int updateWord = getId(word.substr(1, word.size()-1));
                (*sequences[i])[j].first = updateWord;//delete [
                string::size_type eidx;
                unsigned int k = j;//k is temp iterator to prob string
                do
                {
                    k++;
                    tag  = getStr((*sequences[i])[k].second);
                    eidx = tag.find("]");
                } while( eidx == string::npos );
                string nrtag = tag.substr(eidx+1, tag.size()-eidx-1);
                if( nrtag == "i" || nrtag == "l" || nrtag == "nz")
                {
                    continue;
                }
                unsigned int updateTag = getId(nrtag+"b");
                (*sequences[i])[j].second = updateTag;
                j++;
                while( j < k )
                {
                    updateTag = getId(nrtag+"m");
                    (*sequences[i])[j].second = updateTag;
                    j++;
                }
                updateTag = getId(nrtag+"e");
                (*sequences[i])[j].second = updateTag;
                j++;
            }
            else if(tag != "nr" && tag!= "ns" && tag != "nt")
            {
                unsigned int updateTag = getId("o");
                (*sequences[i])[j].second = updateTag;
                j++;
            }
            else
            {
                j++;
            }
        }
    }
}

void Hmm::genSeqs(ostream& ostrm, int seqs)
{
    vector<unsigned long> seq;
    for (int i = 0; i<seqs; i++) {
        genSeq(seq);
        for (unsigned int k = 0; k<seq.size(); k++) {
            if (k)
              ostrm << ' ';
            ostrm << _str2id.getStr(seq[k]);
        }
        ostrm << endl;
        seq.clear();
    }
}


void Hmm::genSeq(vector<unsigned long>& seq)
{
    unsigned long state = _initState, next, obs;
    while (true) {
        if (!_transition.rand(state, next) || !_emission.rand(next, obs))
          break;
        state = next;
        seq.push_back(obs);
    }
}

void PseudoCounts::print(Str2IdMap& str2id)
{
    cerr << "TRANSITION"<<endl;
    _transCount.save(cerr, str2id);
    cerr << "*********************" << endl;
    cerr << "EMISSION"<<endl;
    _emitCount.save(cerr, str2id);
    cerr << "*********************" << endl;
    cerr << "STATE"<<endl;
    _stateCount.save(cerr, str2id);
    cerr << "*********************" << endl;
    cerr << "INIT-PROBS"<<endl;
}
